import requests
from flask import current_app


class OllamaClient:
    def __init__(self, app=None):
        self.app = app

    def init_app(self, app):
        self.app = app

    def generate(self, prompt):
        try:
            response = requests.post(
                f"{current_app.config['OLLAMA_API_BASE']}/chat/completions",
                json={
                    "model": current_app.config['MODEL_NAME'],
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": 0.7
                },
                timeout=300
            )
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except Exception as e:
            current_app.logger.error(f"Ollama API Error: {str(e)}")
            raise


ollama_client = OllamaClient()
generate_response = ollama_client.generate